"""
Compared with ModelContainer classes, ModelController wants to put data close in memory
"""

import torch
from torch import nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR
from torch.utils.tensorboard import SummaryWriter

from path import *
from config import cfg, parser_info
from model import get_models, unfreeze_model
from confusion_matrix import confusion_matrix
from warm_scheduler import get_scheduler
from timer import Timer, TimeAccumulator
from logger import logger
from losses import mse_loss as criterion_mse, FocalLoss

import pandas as pd
import numpy as np
from typing import List
from functools import reduce

timer_batch = Timer('batch')
timer_epoch = Timer('epoch')


def add(x, y):
    return x + y


def pass_epoch(split: str, n_epoch: int, best_n_corrects: list, best_epochs: list, models: List[nn.Module] = None, dataloader: DataLoader = None,
               criteria: list = None, batch_itr: int = 1, epoch_itr: int = 1, **kwargs):
    """
    Train a list of models on the same dataloader.

    models: a list of models
    batch_itr: batch info interval
    epoch_itr: epoch info interval

    save_paths:
        -single model: ['path']
        -multi model: ['path1', 'path2', ..., 'path_ensemble1', 'path_ensemble2', ...]
                      index: i or i + `n_models`

    Return epoch statistics
    """
    assert split in ['train', 'test', 'val']

    # --unfold kwargs-- #
    optimizers = kwargs.get('optimizers', None)  # list of optimizers
    schedulers = kwargs.get('schedulers', None)  # list of schedulers
    save_paths = kwargs.get('save_paths', None)  # list of model paths
    time_accumulator = TimeAccumulator()

    # --basic info-- #
    n_models = len(models)  # number of models
    n_criteria = len(criteria)  # number of criteria
    total_batch = len(dataloader)
    # when train multi-models, save the metrics of ensemble model
    n_ens = 1 if n_models == 1 else n_models + 1

    # --epoch-level statistics-- #
    cms = [torch.zeros((cfg.class_num + 1, cfg.class_num + 1)).int().cuda() for _ in range(n_ens)]
    wrong_path = [[] for _ in range(n_ens)]
    # loss_epoch_scalar = np.zeros((n_ens, n_criteria))  # numbers
    loss_epoch_scalar = [[0 for _ in range(n_criteria)] for _ in range(n_ens)]
    n_corrects = [0 for _ in range(n_ens)]  # accuracy

    # --batch-level statistics-- #
    if split == 'train':
        loss_batch = [torch.zeros(1).cuda() for _ in range(n_models)]  # for back-propagation
    # loss_batch_scalar = np.zeros((n_ens, n_criteria))  # numbers
    loss_batch_scalar = [[0 for _ in range(n_criteria)] for _ in range(n_ens)]

    timer_epoch.clock(event=f'train epoch {n_epoch + 1} starts', verbose=False)

    for model in models:
        if split == 'train':
            model.train()
        else:
            model.eval()

    for batch_num, sample in enumerate(dataloader):
        time_accumulator.set_start(1)
        timer_batch.clock(event=f'{split} epoch {n_epoch + 1}, batch {batch_num + 1} starts', verbose=False)

        # --put data into GPU-- #
        data, label, path = sample
        data = data.cuda()
        label = label.cuda()

        # --feed data into model-- #
        predictions = []
        for model in models:
            prediction = model(data)
            predictions.append(prediction)
        time_accumulator.set_end(1)

        # --calculate loss-- #
        # loss_epoch_scalar[i][j] is the j-th criterion for the i-th model
        time_accumulator.set_start(2)
        if split == 'train':
            loss_batch = [torch.zeros(1).cuda() for _ in range(n_models)]
        for i, prediction in enumerate(predictions):
            for j, criterion in enumerate(criteria):
                loss = criterion(prediction, label)
                if split == 'train':
                    loss_batch[i] += loss
                loss_batch_scalar[i][j] = loss  # .item()
        if n_models > 1:
            for j in range(n_criteria):
                loss_batch_scalar[-1][j] = loss_batch_scalar[0][j]
                for i in range(1, n_models):
                    loss_batch_scalar[-1][j] += loss_batch_scalar[i][j]
            # loss_batch_scalar[-1] = loss_batch_scalar[:-1].sum()
        for j in range(n_criteria):
            for i in range(n_ens):
                loss_epoch_scalar[i][j] += loss_batch_scalar[i][j]
        # loss_epoch_scalar += loss_batch_scalar
        time_accumulator.set_end(2)

        # --back propagation-- #
        time_accumulator.set_start(3)
        if split == 'train':
            for i, optimizer in enumerate(optimizers):
                optimizer.zero_grad()
                loss_batch[i].backward()
                optimizer.step()
        time_accumulator.set_end(3)

        # --calculate metric-- #
        # confusion matrix
        time_accumulator.set_start(4)
        if n_models > 1:
            predictions.append(reduce(add, predictions))
        for i, prediction in enumerate(predictions):
            confusion_matrix(label, prediction, cms[i],
                             img_name=path, wrong_imgs=wrong_path[i])
        time_accumulator.set_end(4)

        # --batch log info-- #
        time_accumulator.set_start(5)
        if (batch_itr > 0) and (0 == (1 + batch_num) % batch_itr):
            timer_batch.clock(event=f'{split} epoch {n_epoch + 1}, batch {batch_num + 1} ends', verbose=True)
            for i in range(n_ens):  # info for every model + ensemble model
                batch_loss_info = '[batch loss: ' + \
                                  '|'.join([f'{loss:.4f}' for loss in loss_batch_scalar[i]]) + \
                                  ']'
                info = f'[batch info]' \
                       f'[net: {i + 1}]' \
                       f'[epoch: {cfg.prefix}{n_epoch + 1}/{cfg.EPOCHS}{cfg.suffix}]' \
                       f'[batch: {batch_num + 1}/{total_batch}]' + \
                       batch_loss_info
                logger.logger.info(info)
        time_accumulator.set_end(5)

    # --change learning rate-- #
    if split == 'train' and schedulers is not None:
        for scheduler in schedulers:
            scheduler.step()

    # --update best performance-- #
    for i in range(n_ens):
        n_corrects[i] = cms[i].diagonal()[:-1].sum().item()
        if n_corrects[i] > best_n_corrects[i]:
            best_n_corrects[i] = n_corrects[i]
            best_epochs[i] = n_epoch + 1
            # save model
            if split != 'train':
                if i < n_models:
                    torch.save(models[i].state_dict(), save_paths[i])
                else:  # use multi-model and is evaluating ensemble result
                    for j in range(n_models):
                        torch.save(models[j].state_dict(), save_paths[j + n_models])

    n_dataset = cms[0][-1, -1].item()

    # --epoch log info-- #
    if (epoch_itr > 0) and (0 == (1 + n_epoch) % epoch_itr):
        timer_epoch.clock(event=f'{split} epoch {n_epoch + 1} ends', verbose=True)
        for i in range(n_ens):
            epoch_loss_info = '[epoch loss: ' + \
                              '|'.join([f'{loss / (batch_num + 1):.4f}' for loss in loss_epoch_scalar[i]]) + \
                              ']'
            info = f'[epoch info]' \
                   f'[net: {i + 1}]' \
                   f'[epoch: {cfg.prefix}{n_epoch + 1}/{cfg.EPOCHS}{cfg.suffix}]' + \
                   epoch_loss_info + \
                   f'[acc: {n_corrects[i] / n_dataset * 100:.2f}%]' \
                   f'[best acc: {best_n_corrects[i] / n_dataset * 100:.2f}% (epoch {best_epochs[i]})]'
            logger.logger.info(info)
            logger.logger.info(cms[i])

    return


if __name__ == '__main__':
    from dataset import BaldClassificationDataset

    if args.multi_model:
        net1, net2 = get_models()
    else:
        net1 = get_models()

    dataset = BaldClassificationDataset(split='test')
    dataloader = DataLoader(dataset, batch_size=args.BATCH, shuffle=args.shuffle)
    # if args.train_ratio < 1:
    #     dataset_val = BaldClassificationDataset(split='val')
    #     dataloader_val = DataLoader(dataset_val, batch_size=args.BATCH)
    # if not args.train_flyai:
    #     dataset_test = BaldClassificationDataset(split='test')
    #     dataloader_test = DataLoader(dataset_test, batch_size=args.BATCH)
    criterion_ce = FocalLoss() if args.use_focal else nn.BCELoss()

    optimizer1 = torch.optim.Adam(params=filter(lambda p: p.requires_grad, net1.parameters()),
                                 lr=cfg.base_lr, betas=(cfg.beta1, cfg.beta2),
                                 eps=cfg.eps, weight_decay=cfg.weight_decay)
    optimizer2 = torch.optim.Adam(params=filter(lambda p: p.requires_grad, net1.parameters()),
                                  lr=cfg.base_lr, betas=(cfg.beta1, cfg.beta2),
                                  eps=cfg.eps, weight_decay=cfg.weight_decay)
    def_scheduler = f'get_scheduler(' \
                    f'warm_epoch=cfg.warm_up, ' \
                    f'max_lr=cfg.lr{1}, ' \
                    f'T_0=cfg.T_0, ' \
                    f'T_mult=cfg.T_mult, ' \
                    f'eta_min=cfg.eta_min{1})'
    def_scheduler = eval(def_scheduler)
    scheduler1 = LambdaLR(optimizer1, def_scheduler)
    def_scheduler = f'get_scheduler(' \
                    f'warm_epoch=cfg.warm_up, ' \
                    f'max_lr=cfg.lr{2}, ' \
                    f'T_0=cfg.T_0, ' \
                    f'T_mult=cfg.T_mult, ' \
                    f'eta_min=cfg.eta_min{2})'
    def_scheduler = eval(def_scheduler)
    scheduler2 = LambdaLR(optimizer2, def_scheduler)

    # train(0, [0, 0, 0], [0, 0, 0], [net1, net2], dataloader, [optimizer1, optimizer2], [criterion_ce], 1, 1,
    #       schedulers=[scheduler1, scheduler2])
    with torch.no_grad():
        pass_epoch('test', 0, [0, 0, 0], [0, 0, 0], [net1, net2], dataloader, [criterion_ce], 0, 1,
                   optimizers=[optimizer1, optimizer2], schedulers=[scheduler1, scheduler2],
                   save_paths=[MODEL_FILE_PATH12, MODEL_FILE_PATH22, MODEL_FILE_PATH11, MODEL_FILE_PATH21])